package cn.yhq.http.core;

import java.io.IOException;

import okhttp3.MediaType;
import okhttp3.MultipartBody;
import okhttp3.ResponseBody;
import okio.Buffer;
import okio.BufferedSource;
import okio.ForwardingSource;
import okio.Okio;
import okio.Source;


public class ProgressResponseBody extends ResponseBody {
  private final ResponseBody responseBody;
  private final ProgressListener progressListener;
  private BufferedSource bufferedSource;

  public ProgressResponseBody(ResponseBody responseBody, ProgressListener progressListener) {
    this.responseBody = responseBody;
    this.progressListener = progressListener;
  }

  @Override
  public MediaType contentType() {
    return responseBody.contentType();
  }

  @Override
  public long contentLength() {
    return responseBody.contentLength();
  }

  @Override
  public BufferedSource source() {
    if (bufferedSource == null) {
      bufferedSource = Okio.buffer(source(responseBody.source()));
    }
    return bufferedSource;
  }

  private Source source(Source source) {
    return new ForwardingSource(source) {
      long totalBytesRead = 0L;

      @Override
      public long read(Buffer sink, long byteCount) throws IOException {
        long bytesRead = super.read(sink, byteCount);
        // read() returns the number of bytes read, or -1 if this source is exhausted.
        totalBytesRead += bytesRead != -1 ? bytesRead : 0;
        if (responseBody.contentType() != null) {
          progressListener.update(
              responseBody.contentType().toString().contains(MultipartBody.FORM.toString()),
              totalBytesRead, responseBody.contentLength(), bytesRead == -1);
        } else {
          progressListener.update(false, totalBytesRead, responseBody.contentLength(),
              bytesRead == -1);
        }

        return bytesRead;
      }
    };
  }
}
